# -*- coding: utf8 -*-

import zlib, bz2, lzma
from project.mysql.table import Table
from pymysql import Binary

class TableImageVectorBitstr( Table ):

    _table_name = "image_vector_bitstr"

    def create( self ):
        cur = self.getCursor()

        # Create table
        cur.execute(
            "CREATE TABLE IF NOT EXISTS %s (" % self._table_name+
            "image_id INTEGER UNSIGNED PRIMARY KEY, "+
            "file_id INTEGER UNSIGNED NOT NULL, "+
            "bitmask BIGINT UNSIGNED NOT NULL, "+
            "mode CHAR(1) NOT NULL, "+
            "vector BLOB NOT NULL"+
            ") ENGINE %s" % self.getEngine()
        )
        try:
            cur.execute(
                "ALTER TABLE `%s` ADD INDEX bitmask (bitmask)" % self._table_name
            )
        except:
            pass
        try:
            cur.execute(
                "ALTER TABLE `%s` ADD INDEX files (file_id)" % self._table_name
            )
        except:
            pass
        cur.close()
        return


    def insert( self, file_id, image_id, bitmask, vector ):

        v_str = vector

        len_array = len(v_str)

        gz_str = zlib.compress(v_str, 4)
        len_gz = len(gz_str)

        bz_str = bz2.compress(v_str, 2)
        len_bz = len(bz_str)

        lz_str = lzma.compress(v_str, 2)
        len_lz = len(lz_str)

        code = 's'
        if len_array > len_gz:
            code = 'g'
            v_str = gz_str
            len_array = len_gz
        if len_array > len_bz:
            code = 'b'
            v_str = bz_str
            len_array = len_bz
        if len_array > len_lz:
            code = 'l'
            v_str = lz_str
            len_array = len_lz

        v_str = Binary(v_str)
        cur = self.getCursor()
        v_id = cur.execute(
            "INSERT INTO %s(" % self._table_name+
            "file_id, image_id, bitmask, mode, vector"
            ") VALUES (%s, %s, %s, %s, %s)", (
                file_id, image_id, bitmask, code, v_str
            )
        )
        self.getConnection().commit()
        cur.close()
        return v_id


    def get( self, image_id ):
        cur = self.getCursor()
        cur.execute(
            "SELECT * FROM %s " % self._table_name+
            "WHERE image_id=%(id)s",
            {"id": image_id}
        )
        vector = cur.fetchone()
        cur.close()

        code = vector["mode"]

        v_str = vector["vector"]

        if code == 'g':
            v_str = zlib.decompress(v_str)
        if code == 'b':
            v_str = bz2.decompress(v_str)
        if code == 'l':
            v_str = lzma.decompress(v_str)

        vector["vector"] = v_str

        return vector

    def get_equal_bitmask( self, image_id, bitmask ):
        """
        @return image ids, bitmask which equals to image_id's
        """
        cur = self.getCursor()
        cur.execute(
            "SELECT image_id FROM %s " % self._table_name+
            "WHERE bitmask=%(mask)s AND image_id!=%(id)s",
            {"id": image_id, "mask":bitmask}
        )
        vectors = cur.fetchall()
        cur.close()
        return vectors


    def remove_by_file(self, file_id):
        cur = self.getCursor()
        rc = cur.execute("DELETE FROM %s " % self._table_name +
                         "WHERE file_id=%(id)s", {"id": file_id})
        self.getConnection().commit()
        cur.close()
        return rc


    pass
